import torch  # pytor库，必用
import torch.nn as nn  # 导入nn，这个库里面是一些全连接网络，池化层，这种神经网络中非常基本的一些模块，这里我们主要是用nn.linear是一个全连接层
from transformers import BertModel, BertTokenizer  # transformer库是一个把各种预训练模型集成在一起的库，导入之后，你就可以选择性的使用自己想用的模型，这里使用的BERT模型。所以导入了bert模型，和bert的分词器，这里是对bert的使用，而不是bert自身的源码，如果有时间或者兴趣的话我会在另一篇文章写bert的源码实现


class BertClassification(nn.Module):  # 括号里面是继承什么，该类继承自nn.module，因为我们的模型在根本上是神经网络的，所以继承nn.modelu，继承它的基本属性是自然的了
    def __init__(self, device, class_num):
        super(BertClassification, self).__init__()
        self.model_name = 'hfl/chinese-bert-wwm'
        self.model = BertModel.from_pretrained(self.model_name).to(device)
        self.tokenizer = BertTokenizer.from_pretrained(self.model_name)
        self.fc = nn.Linear(768, class_num).to(device)  # 768取决于BERT结构，2-layer, 768-hidden, 12-heads, 110M parameters
        self.device = device

    def forward(self, x):  # 这里的输入x是一个list,也就是输入文本x：RNG全队梦游失误频频不敌FW，后续淘汰赛成绩引人担忧，我这里是用一句话举例子，实际上的数据是很多很多句话（哈哈，好不专业，很多很多）
        batch_tokenized = self.tokenizer.batch_encode_plus(x, add_special_tokens=True, max_length=150,
                                                           pad_to_max_length=True)  # tokenize、add special token、pad
        input_ids = torch.tensor(batch_tokenized['input_ids']).to(self.device)
        attention_mask = torch.tensor(batch_tokenized['attention_mask']).to(self.device)
        hiden_outputs = self.model(input_ids, attention_mask=attention_mask)
        outputs = hiden_outputs[0][:, 0, :]  # [0]表示输出结果部分，[:,0,:]表示[CLS]对应的结果
        output = self.fc(outputs)
        return output



